{-# LANGUAGE BangPatterns, TypeSynonymInstances, FlexibleInstances  #-}

module Sim.PSF (
    PsfData(..), PSF,  BBox, Pixel,
    toPsf,
    gaussPsf
  ) where

import        GHC.Float   ( int2Double )

import Data.VectorSpace   ( (^-^), (^+^), (*^), (^/) )
import Data.Array         ( Array, (!), array, elems, range)
import Data.Array.Unboxed ( UArray, amap )
import Data.List          ( foldl' )

import Interpolation      ( biLerp )
import Sim.Helper         ( (<~), memo2, pair, floor' )

import qualified Data.Map as M

type Z = Int
type R = Double

type Zero       = (R,R)
type Step       = R
type ScaleRatio = Z
type Pixel      = ((Z,Z), R)
type Kernel     = (Z,Z) -> [(Z,Z)]

type Grid    = (Zero, Step)
type BBox    = ((Z,Z), (Z,Z))
type Nodes   = Array (Z,Z) ((Z,Z) -> R)

type PSF     = (R,R) -> [Pixel]

data PsfData = PsfData Grid BBox Nodes

class Scalable a where
  toScale :: a -> ScaleRatio -> a

instance Scalable Grid where
  toScale (zp,h) r = (zp',h')
     where
        zp' = zp ^+^ (h',h') ^/ 2 -- Зачем нужно второе слагаемое?
        h'  = h / fromIntegral r

instance Scalable BBox where
  toScale (lb,rb) r = (lb',rb')
     where
        lb' = pair (`div` r) lb
        rb' = pair (`div` r) rb

toPsf :: PsfData -> ScaleRatio -> Grid -> ((R,R), (R,R)) -> [Pixel]
toPsf (PsfData mgrid bbox nodes) sratio pgrid = pixPsf psf sbox pgrid
  where
     psf  = interpolate mgrid psf'

     psf' = (<~) $ amap (\f -> f `toProjectWith` sratio `toStretchOn` sgrid) nodes

     sgrid = pgrid `toScale` sratio  -- субпиксельная сетка
     sbox  = bbox  `toScale` sratio

{--
gaussPsf :: (R,R) -> ScaleRatio -> Grid -> ((R,R), (R,R)) -> [Pixel]
gaussPsf s@(sx,sy) sratio pgrid = pixPsf psf (l,r) pgrid
  where
     psf   = gauss `toProjectWith` sratio `toStretchOn` sgrid

     sgrid = pgrid `toScale` sratio -- субпиксельная сетка

     l     = pair round $ (-6) *^ s
     r     = pair round $ ( 6) *^ s

     gauss :: (Z,Z) -> R
     gauss (x,y) = exp( -0.5*( a*a + b*b )) / ( 2 * pi * sx * sy ) 
       where
          a = (int2Double x) / sx
          b = (int2Double y) / sy
--}
--{--
type Scale  = Z
type Width  = Z
type Height = Z
type Image  = ((Width,Height), [Pixel])

gaussPsf :: (R,R) -> Scale -> Grid -> ((R,R), (R,R)) -> [Pixel]
gaussPsf s@(sx,sy) n = pixels (templates img n) n
  where
    img :: Image
    img = (wh,pImg)

    wh :: (Width,Height)
    wh = pair round $ 10 *^ s

    pImg :: [Pixel]
    pImg  = [ (i+c, gauss i) | i <- range (-c,c) ]
      where
       c = pair (`div` 2) wh 

    gauss :: (Z,Z) -> R
    gauss (x,y) = exp( -0.5*( a*a + b*b )) / ( 2 * pi * sx * sy )
      where
        a = (int2Double x) / sx
        b = (int2Double y) / sy

templates :: Image -> Scale -> Array (Z,Z) [Pixel] 
templates (wh,img) n = a
  where
    b = (0,(n,n)-1)
    c = pair (`div` 2) wh
    a = array b [(i, accum [(pair (`div` n) (j+i-c), v) | (j,v) <- img]) | i <- range b]

pixels :: Array (Z,Z) [Pixel] -> Scale -> Grid -> ((R,R), (R,R)) -> [Pixel]
pixels a n (zp,h) (_,xy)  = [ (p+i,v) | (p,v) <- a ! pair (`mod` n) j]
  where
    i = pair floor' $ (xy - zp) ^/h   
    j = pair floor' $ (xy - zp) ^/h ^/(fromIntegral n)

accum :: (Num v, Ord k) => [(k,v)] -> [(k,v)]
accum xs = M.assocs $ foldr f M.empty xs
  where
    f (k,v) m = case k `M.lookup` m of
                  Just e  -> M.insert k (e+v) m
                  Nothing -> M.insert k v     m
--}


{--
pixPsf :: ((R,R) -> (Z,Z) -> R) -> BBox -> Grid -> (R,R) -> [Pixel]
pixPsf psf box gr xy = [ (ci ^+^ i, v) | (i, v) <- pixelateIn box (psf xy) ]
  where
    ci = indexOn gr xy
--}
pixPsf :: ((R,R) -> (Z,Z) -> R) -> BBox -> Grid -> ((R,R), (R,R)) -> [Pixel]
pixPsf psf box gr (uv,xy) = [ (ci ^+^ i, v) | (i, v) <- pixelateIn box (psf uv) ]
  where
    ci = indexOn gr xy

interpolate :: Grid -> ((Z,Z) -> (R,R) -> (Z,Z) -> R) -> (R,R) -> (Z,Z) -> R
interpolate (zp, h) f p = f' p' p
  where
    f' = (uncurry . biLerp . curry) f

    p' = (p ^-^ zp) ^/ h

toProjectWith ::((Z,Z) -> R) -> ScaleRatio -> (Z,Z) -> (Z,Z) -> R
toProjectWith f r = proj
  where
     proj = ( memo2 cs $ сonv k . translate f ) . pair (`mod` r)

     -- определяем ядро свертки
     k (i,j) = range (lb,rb)
       where
          lb  = (i*r - hr, j*r - hr)
          rb  = (i*r + hr, j*r + hr)

     -- размер кэша для мемозации.
     cs = ((-120,-120), (120,120)) 

     hr = r `div` 2

-- Свертка с заданным ядром функции, определенной на множестве пар целых
-- чисел.
сonv :: Kernel -> ((Z,Z) -> R) -> (Z,Z) -> R
сonv k f ij = foldl' (+) 0 [ f x | x <- k ij ]

-- Параллельный перенос и на заданный вектор функции.
translate :: ((Z,Z) -> R) -> (Z,Z) -> (Z,Z) -> R 
translate f dxy xy = f $ xy ^-^ dxy

-- Операция "натягивания" функции, определенной на множестве пар целых
-- чисел, на множество пар действительных чисел.
toStretchOn :: ((Z,Z) -> b) -> Grid-> (R,R) -> b
toStretchOn f gr  =  f . indexOn gr

-- 
indexOn :: Grid -> (R,R) -> (Z,Z)
indexOn (zp,h) xy = ij
 where
   ij = pair floor' $ (xy ^-^ zp) ^/h ^+^ (0.5,0.5)


-- Пикселизация функции в заданном прямоугольнике.
pixelateIn :: BBox -> ((Z,Z) -> R) -> [Pixel]
pixelateIn b f = [ (ij,  f ij) | ij <- range b ]
